import torch.nn.functional as F
from PIL import Image
import numpy as np
import torchvision.transforms as transforms
from torchvision import models
from net import FCN8s
import torch


class PredictUtil:
    def __init__(self, modelSavePath: str, net: FCN8s, colormap):
        net.load_state_dict(torch.load(modelSavePath))
        net.eval()
        self.net = net.cuda()
        self.colorMap = np.array(colormap).astype('uint8')
        self.ImageTransToTensor = transforms.Compose([transforms.ToTensor(),
                                                      transforms.Normalize([0.485, 0.456, 0.406],
                                                                           [0.229, 0.224, 0.225])])

    def __call__(self, predictImagePath: str, resultImageSavePath: str):
        print('*' * 10, 'begin predict')
        preImage = Image.open(predictImagePath)
        preImage = self.ImageTransToTensor(preImage).cuda()  # torch.Size([3, 352, 352])
        preImage = torch.unsqueeze(preImage, dim=0).cuda()  # torch.Size([1, 3, 352, 352])

        netOutScoresImage = self.net(preImage).cuda()

        netOutRateImage = F.log_softmax(netOutScoresImage, dim=1)

        predicMatrix = netOutRateImage.max(1)[1].squeeze().cpu().data.numpy()
        resultPicArray = self.colorMap[predicMatrix]

        resultImage = Image.fromarray(resultPicArray).convert('RGBA')

        resultImage.save(resultImageSavePath)
        resultImage.show()


from PIL.Image import NEAREST, BILINEAR, BICUBIC, LANCZOS, BOX, HAMMING

# begin train

# 预训练对象
pretrained_net = models.vgg16_bn(pretrained=False)
# 模型保存路径
modelSavePath = 'epoch_0_checkpoint.pt'
# 创建网络对象
net = FCN8s(pretrained_net, 5)
# 可视化结果的Map
# 下面是 RGBA 的值
colormap = [[0, 0, 0, 0], [0, 255, 127, 200], [255, 255, 0, 200], [0, 191, 255, 200], [220, 20, 60, 200]]

# 预测图片的路径
predictImagePath = r'data/test/test01.jpg'
# 结果保存的路径以及名称
predictResultSavePath = r'data/test/testing2-1.png'

# begin to predict
predictUtil = PredictUtil(modelSavePath, net, colormap)
predictUtil(predictImagePath, predictResultSavePath)
